import sys
import cv2
import numpy as np
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QHBoxLayout, QPushButton, QFileDialog, QComboBox, QLabel, QGroupBox, QSizePolicy
from PyQt5.QtGui import QPixmap, QImage
import os
from Function import *


class ImageProcessingApp(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()
        self.image = None
        self.processed_image = None
        # 获取当前脚本所在的目录
        self.current_dir = os.path.dirname(os.path.abspath(__file__))

    def initUI(self):
        # 创建主布局
        main_layout = QVBoxLayout()
        main_layout.setContentsMargins(20, 20, 20, 20)
        main_layout.setSpacing(20)

        # 创建按钮和下拉框布局
        button_layout = QHBoxLayout()
        button_layout.setSpacing(10)

        # 创建按钮
        self.load_button = QPushButton('加载图像', self)
        self.process_button = QPushButton('处理图像', self)
        self.save_button = QPushButton('保存处理后的图像', self)
        self.delete_button = QPushButton('删除处理后的图像', self)
        self.load_button.clicked.connect(self.load_image)
        self.process_button.clicked.connect(self.process_image)
        self.save_button.clicked.connect(self.save_image)
        self.delete_button.clicked.connect(self.delete_processed_image)
        self.save_button.setEnabled(False)
        self.delete_button.setEnabled(False)

        # 设置按钮样式
        button_style = """
            QPushButton {
                background-color: #4CAF50;
                color: white;
                padding: 10px 20px;
                border: none;
                border-radius: 5px;
                font-size: 14px;
            }
            QPushButton:hover {
                background-color: #45a049;
            }
        """
        self.load_button.setStyleSheet(button_style)
        self.process_button.setStyleSheet(button_style)
        self.save_button.setStyleSheet(button_style)
        self.delete_button.setStyleSheet(button_style)

        button_layout.addWidget(self.load_button)
        button_layout.addWidget(self.process_button)
        button_layout.addWidget(self.save_button)
        button_layout.addWidget(self.delete_button)

        # 创建下拉框
        self.algorithm_combobox = QComboBox(self)
        self.algorithm_combobox.addItems([
            '线性变换', '对数变换', '幂律变换', '分段线性变换', '直方图均衡化', '直方图规定化',
            '均值滤波', '高斯滤波', '中值滤波', '梯度算子', '拉普拉斯算子', '高通滤波',
            'Sobel边缘检测', 'Prewitt边缘检测', '傅里叶变换', '离散余弦变换',
            '理想低通滤波', '巴特沃斯低通滤波', '理想高通滤波', '拉普拉斯频域滤波',
            '带通滤波', '带阻滤波', '逆滤波', '维纳滤波', '约束最小二乘滤波',
            '迭代盲复原', '凸集投影', '腐蚀', '膨胀', '开运算', '闭运算',
            '顶帽变换', '黑帽变换', '形态学梯度', '骨架提取', '连通域分析',
            '平移', '旋转', '近邻插值缩放', '双线性插值缩放', '镜像', '错切',
            '透视变换', '仿射变换', '校正几何畸变', 'SIFT图像配准', 'ORB图像配准'
        ])
        # 设置下拉框样式
        combobox_style = """
            QComboBox {
                padding: 10px;
                border: 1px solid #ccc;
                border-radius: 5px;
                font-size: 14px;
            }
        """
        self.algorithm_combobox.setStyleSheet(combobox_style)
        button_layout.addWidget(self.algorithm_combobox)

        # 创建图像显示布局
        image_layout = QHBoxLayout()
        image_layout.setSpacing(20)

        # 创建原始图像显示组
        original_group = QGroupBox('原始图像')
        original_layout = QVBoxLayout()
        self.original_image_label = QLabel(self)
        self.original_image_label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        self.original_image_label.setAlignment(Qt.AlignCenter)
        original_layout.addWidget(self.original_image_label)
        original_group.setLayout(original_layout)
        # 设置组框样式
        groupbox_style = """
            QGroupBox {
                border: 1px solid #ccc;
                border-radius: 5px;
                padding: 10px;
                font-size: 14px;
            }
            QGroupBox::title {
                subcontrol-origin: margin;
                left: 10px;
                padding: 0 5px;
            }
        """
        original_group.setStyleSheet(groupbox_style)

        # 创建处理后图像显示组
        processed_group = QGroupBox('处理后图像')
        processed_layout = QVBoxLayout()
        self.processed_image_label = QLabel(self)
        self.processed_image_label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        self.processed_image_label.setAlignment(Qt.AlignCenter)
        processed_layout.addWidget(self.processed_image_label)
        processed_group.setLayout(processed_layout)
        processed_group.setStyleSheet(groupbox_style)

        image_layout.addWidget(original_group)
        image_layout.addWidget(processed_group)

        main_layout.addLayout(button_layout)
        main_layout.addLayout(image_layout)

        self.setLayout(main_layout)
        self.setWindowTitle('图像处理应用')
        self.setGeometry(300, 300, 1000, 600)
        # 设置窗口背景颜色
        self.setStyleSheet("background-color: #f0f0f0;")
        self.show()

    def load_image(self):
        file_dialog = QFileDialog()
        file_path, _ = file_dialog.getOpenFileName(self, '选择图像文件', '', '图像文件 (*.png *.jpg *.bmp)')
        if file_path:
            self.image = cv2.imread(file_path, cv2.IMREAD_GRAYSCALE)
            if self.image is not None:
                self.display_image(self.original_image_label, self.image)
                self.processed_image = None
                self.processed_image_label.clear()
                self.save_button.setEnabled(False)
                self.delete_button.setEnabled(False)
            else:
                print(f"无法读取文件: {file_path}，请检查文件路径和完整性。")

    def process_image(self):
        if self.image is None:
            print("请先加载图像。")
            return
        algorithm = self.algorithm_combobox.currentText()
        result = None
        if algorithm == '线性变换':
            result = linear_transformation(self.image, a=1.5, b=30)
        elif algorithm == '对数变换':
            result = log_transformation(self.image, c=1)
        elif algorithm == '幂律变换':
            result = gamma_transformation(self.image, gamma=0.5)
        elif algorithm == '分段线性变换':
            result = piecewise_linear_transformation(self.image, r1=50, s1=20, r2=200, s2=230)
        elif algorithm == '直方图均衡化':
            result = histogram_equalization(self.image)
        elif algorithm == '直方图规定化':
            reference_path = os.path.join(self.current_dir, 'girl.bmp')
            reference = cv2.imread(reference_path, cv2.IMREAD_GRAYSCALE)
            if reference is None:
                print(f"无法读取文件: {reference_path}，请检查文件路径和完整性。")
            else:
                result = histogram_matching(self.image, reference)
        elif algorithm == '均值滤波':
            result = mean_filter(self.image, kernel_size=3)
        elif algorithm == '高斯滤波':
            result = gaussian_filter(self.image, kernel_size=3, sigma=0)
        elif algorithm == '中值滤波':
            result = median_filter(self.image, kernel_size=3)
        elif algorithm == '梯度算子':
            result = gradient_operator(self.image)
        elif algorithm == '拉普拉斯算子':
            result = laplacian_operator(self.image)
        elif algorithm == '高通滤波':
            result = high_pass_filter(self.image)
        elif algorithm == 'Sobel边缘检测':
            result = sobel_edge_detection(self.image)
        elif algorithm == 'Prewitt边缘检测':
            result = prewitt_edge_detection(self.image)
        elif algorithm == '傅里叶变换':
            result = fft_transform(self.image)
        elif algorithm == '离散余弦变换':
            result = dct_transform(self.image)
        elif algorithm == '理想低通滤波':
            result = ideal_low_pass_filter(self.image, cutoff=30)
        elif algorithm == '巴特沃斯低通滤波':
            result = butterworth_low_pass_filter(self.image, cutoff=30, order=2)
        elif algorithm == '理想高通滤波':
            result = ideal_high_pass_filter(self.image, cutoff=30)
        elif algorithm == '拉普拉斯频域滤波':
            result = laplacian_frequency_filter(self.image)
        elif algorithm == '带通滤波':
            result = band_pass_filter(self.image, low_cutoff=10, high_cutoff=30)
        elif algorithm == '带阻滤波':
            result = band_stop_filter(self.image, low_cutoff=10, high_cutoff=30)
        elif algorithm == '逆滤波':
            kernel = np.ones((3, 3), np.float32) / 9
            result = inverse_filter(self.image, kernel, noise_power=0)
        elif algorithm == '维纳滤波':
            kernel = np.ones((3, 3), np.float32) / 9
            result = wiener_filter(self.image, kernel, K=0.01)
        elif algorithm == '约束最小二乘滤波':
            kernel = np.ones((3, 3), np.float32) / 9
            result = constrained_least_squares_filter(self.image, kernel, gamma=0.01)
        elif algorithm == '迭代盲复原':
            result = iterative_blind_deconvolution(self.image, kernel_size=3, iterations=3)
        elif algorithm == '凸集投影':
            kernel = np.ones((3, 3), np.float32) / 9
            result = projection_onto_convex_sets(self.image, kernel, iterations=10)
        elif algorithm == '腐蚀':
            result = erosion(self.image, kernel_size=3)
        elif algorithm == '膨胀':
            result = dilation(self.image, kernel_size=3)
        elif algorithm == '开运算':
            result = opening(self.image, kernel_size=3)
        elif algorithm == '闭运算':
            result = closing(self.image, kernel_size=3)
        elif algorithm == '顶帽变换':
            result = top_hat(self.image, kernel_size=3)
        elif algorithm == '黑帽变换':
            result = black_hat(self.image, kernel_size=3)
        elif algorithm == '形态学梯度':
            result = morphological_gradient(self.image, kernel_size=3)
        elif algorithm == '骨架提取':
            result = skeleton_extraction(self.image)
        elif algorithm == '连通域分析':
            num_labels, labels, stats, centroids = connected_components_analysis(self.image)
            result = np.uint8(labels * (255 / num_labels))
        elif algorithm == '平移':
            result = translation(self.image, x=50, y=50)
        elif algorithm == '旋转':
            result = rotation(self.image, angle=45, scale=1.0)
        elif algorithm == '近邻插值缩放':
            result = nearest_neighbor_scaling(self.image, scale_x=1.5, scale_y=1.5)
        elif algorithm == '双线性插值缩放':
            result = bilinear_scaling(self.image, scale_x=1.5, scale_y=1.5)
        elif algorithm == '镜像':
            result = mirror(self.image, axis=1)
        elif algorithm == '错切':
            result = shear(self.image, shear_factor_x=0.2, shear_factor_y=0.2)
        elif algorithm == '透视变换':
            src_points = np.float32([[0, 0], [self.image.shape[1], 0], [0, self.image.shape[0]], [self.image.shape[1], self.image.shape[0]]])
            dst_points = np.float32([[50, 50], [self.image.shape[1] - 50, 50], [50, self.image.shape[0] - 50], [self.image.shape[1] - 50, self.image.shape[0] - 50]])
            result = perspective_transformation(self.image, src_points, dst_points)
        elif algorithm == '仿射变换':
            src_points = np.float32([[0, 0], [self.image.shape[1], 0], [0, self.image.shape[0]]])
            dst_points = np.float32([[50, 50], [self.image.shape[1] - 50, 50], [50, self.image.shape[0] - 50]])
            result = affine_transformation(self.image, src_points, dst_points)
        elif algorithm == '校正几何畸变':
            src_points = np.float32([[0, 0], [self.image.shape[1], 0], [0, self.image.shape[0]], [self.image.shape[1], self.image.shape[0]]])
            dst_points = np.float32([[50, 50], [self.image.shape[1] - 50, 50], [50, self.image.shape[0] - 50], [self.image.shape[1] - 50, self.image.shape[0] - 50]])
            result = geometric_distortion_correction(self.image, src_points, dst_points, degree=2)
        elif algorithm == 'SIFT图像配准':
            reference_path = os.path.join(self.current_dir, 'girl.bmp')
            reference = cv2.imread(reference_path, cv2.IMREAD_GRAYSCALE)
            if reference is None:
                print(f"无法读取文件: {reference_path}，请检查文件路径和完整性。")
            else:
                result = sift_image_registration(self.image, reference)
        elif algorithm == 'ORB图像配准':
            reference_path = os.path.join(self.current_dir, 'girl.bmp')
            reference = cv2.imread(reference_path, cv2.IMREAD_GRAYSCALE)
            if reference is None:
                print(f"无法读取文件: {reference_path}，请检查文件路径和完整性。")
            else:
                result = orb_image_registration(self.image, reference)

        if result is not None:
            self.processed_image = result
            self.display_image(self.processed_image_label, result)
            self.save_button.setEnabled(True)
            self.delete_button.setEnabled(True)
        else:
            print(f"图像处理失败，请检查算法实现或输入图像。")

    def save_image(self):
        if self.processed_image is not None:
            file_dialog = QFileDialog()
            file_path, _ = file_dialog.getSaveFileName(self, '保存图像', '', '图像文件 (*.png *.jpg *.bmp)')
            if file_path:
                cv2.imwrite(file_path, self.processed_image)
        else:
            print("没有处理后的图像可保存。")

    def delete_processed_image(self):
        self.processed_image = None
        self.processed_image_label.clear()
        self.save_button.setEnabled(False)
        self.delete_button.setEnabled(False)

    def display_image(self, label, image):
        if len(image.shape) == 2:  # 灰度图像
            height, width = image.shape
            qimage = QImage(image.data, width, height, QImage.Format_Grayscale8)
        else:  # 彩色图像
            height, width, channel = image.shape
            bytes_per_line = 3 * width
            qimage = QImage(image.data, width, height, bytes_per_line, QImage.Format_RGB888).rgbSwapped()
        pixmap = QPixmap.fromImage(qimage)
        label.setPixmap(pixmap.scaled(label.width(), label.height(), aspectRatioMode=Qt.KeepAspectRatio))


if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = ImageProcessingApp()
    sys.exit(app.exec_())